import numpy as np
import cv2
import os

def img2vector(imgfilename):
    img = cv2.imread(imgfilename, cv2.IMREAD_GRAYSCALE)
    rows, columns = img.shape
    #print(img.shape)
    vec = img.reshape(rows * columns)
    return vec


training_dir = "../training"
sub_dir_and_files = os.listdir(training_dir)

sub_dirs = []
# 如果是目录
for x in sub_dir_and_files:
    if os.path.isdir(training_dir + "/" + x):
        sub_dirs.append(x)

# 计算装备图像总数
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(training_dir + "/" + subdir))

# 初始化训练图像数据矩阵 (N 行，128*128列) 和装备向量（长度为N）
training_img_matrix = np.zeros((N, 128 * 128))  # 每个图像一行数据
training_equipment_vector = [''] * N

i = 0  # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(training_dir + "/" + subdir)
    for image in image_files:
        # 将图像转换为向量
        v = img2vector(training_dir + "/" + subdir + "/" + image)
        training_img_matrix[i] = v
        training_equipment_vector[i] = subdir
        i += 1


# 获取测试目录下的所有子目录
testing_dir = '../testing'
sub_dir_and_files = os.listdir(testing_dir)
sub_dirs = []
for x in sub_dir_and_files:
    if os.path.isdir(testing_dir + '/' + x):
        sub_dirs.append(x)


# 计算测试目录下所有装备图像的数量和
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(testing_dir + '/' + subdir))

# 初始化测试图像数据矩阵（N行，128*128列）和（相对应的）装备向量（长度为N）
testing_image_matrix = np.zeros((N, 128*128)) # 每个测试图像对应一行所有像素值
testing_equipment_vector = [""]*N # 与同下标位置测试图像相对应的装备名称


# 将三个子目录下所有图像文件读入图像数据矩阵并标识装备类别名称
i = 0 # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(testing_dir + "/" + subdir)
    for image in image_files:
        # 图像不在当前目录，因此需要加上目录前缀
        v = img2vector(testing_dir+"/"+subdir+"/"+image)
        testing_image_matrix[i] = v
        testing_equipment_vector[i] = subdir
        i += 1

# 导入 sklearn
from sklearn import neighbors

# 取得 knn 分类器  n_neighbors=8 表示设置 K 值
knn = neighbors.KNeighborsClassifier(n_neighbors=5)

# 特征向量
data = training_img_matrix

# 标记数据 1 为 喜剧片， 2 为爱情片 ， 3 为 动作片
labels = training_equipment_vector

# 训练模型
fit = knn.fit(data, labels)

# 预测
result = fit.predict(testing_image_matrix)

# 输出预测结果
print(result)

# 预测正确的数量
count = 0
for i in range(len(result)):
    if result[i] == testing_equipment_vector[i]:
        count += 1

print( '%.1f%%' %(count / len(result) *100))

from sklearn.metrics import accuracy_score

print(accuracy_score(result, testing_equipment_vector))

# 获最近的 K 个邻居
# movie = np.array([[1,9,1]])
# 默认K值为5 n_neighbors = 4 表示设置 K = 4
# 返回前 4 个邻居的“索引”位置
# neighbors=fit.kneighbors(movie,return_distance=False, n_neighbors = 4)

# 输出所有的邻居的索引位置
# print(neighbors)
# 构建Knn 分类器
